/*
  class for copmutation of stream function of 2D velocity field
*/
#include "fem.h"
#include "output.h"
#include "integration/quadrature.h"
#include "stateproblem.h"

class C3DFlowProblem : public CStateProblem
{
 public:
  virtual void afterReadConfig();
  void calculateToF();
};

class C3DToFProblem : public CStateProblem
{
 private:
  int v_type;
  IntPoint<3> *vip3d;
  IntPoint<2> *vip2d;
  IntPoint<1> *vip1d;
  CVec *vq;
  CVar vx, vy, vz;
 public:
  void afterLoadMesh();
  void afterEvaluateVolumeComponents(int elem, int qp);
  void setVelocity(int type, CVec *_vq, IntPoint<3> *ip3d, IntPoint<2> *ip2d, IntPoint<1> *ip1d);
  void setMesh(Mesh *m, SaveCells cells);
  bool useVolumeIntegral() { return true; };
  bool useFaceIntegral() { return false; };
  bool useEdgeIntegral() { return false; };
  CVar volumeIntegral(int comp, int elem, int qp);
  
  friend class C3DFlowProblem;
};


C3DToFProblem SFP;


void C3DFlowProblem::afterReadConfig()
{
  SaveCells c = (output->saveCells(Save3D)?Save3D:SaveNone)
              | (output->saveCells(Save2D)?Save2D:SaveNone)
              | (output->saveCells(Save1D)?Save1D:SaveNone);
  SFP.setMesh(mesh, c);
  SFP.registerQuadratures(Q3d, Q2d, Q1d);
  SFP.setVelocity(FE_LIN, &q, U[0]->ip3d, U[0]->ip2d, U[0]->ip1d);
  SFP.params["savebase"] = params["savebase"] + "_tof";
  
  string bctag;
  for (map<int,string>::iterator bcname=mesh->beginPhysicalNames(); bcname!=mesh->endPhysicalNames(); bcname++)
  {
    bctag = "bc[" + bcname->second + ",tof]";
    if (params.find(bctag) != params.end())
      SFP.params["bc[" + bcname->second + ",0]"] = params[bctag];
  }
  SFP.addScalar("tof", FE_LIN);
}

void C3DFlowProblem::calculateToF()
{
  double xnorm;
  log(W_PROG, "\nCalculating time of flight...\n------------------------------\n");
  SFP.prepare();
  SFP.assemble();
  SFP.solve(xnorm, 1.0);
  SFP.saveVTKview();
  SFP.saveSolution();
}




void C3DToFProblem::afterLoadMesh()
{
	mesh->findChildrenParents(true);
}


void C3DToFProblem::setVelocity(int type, CVec *_vq, IntPoint<3> *ip3d, IntPoint<2> *ip2d, IntPoint<1> *ip1d)
{
	v_type = type;
	vq = _vq;
	vip3d = ip3d;
	vip2d = ip2d;
	vip1d = ip1d;
};


void C3DToFProblem::afterEvaluateVolumeComponents(int elem, int qp)
{
//  CVec g;
//  vip3d->calcE(elem, EVAL_FG);
//  vip3d->calcQP(qp, EVAL_FG);
//  vip3d->eval(*vq, elem, 0, vx, g, EVAL_FG);
//  vip3d->eval(*vq, elem, vip3d->getMeshNdof(), vy, g, EVAL_FG);
//  vip3d->eval(*vq, elem, 2*vip3d->getMeshNdof(), vz, g, EVAL_FG);

	int ind[4], offset;

  	for (int i=0; i<4; i++)
  		ind[i] = vip3d->mapDof(elem, i);

  	offset = vip3d->getMeshNdof();

  	vx = (*vq)(ind[0])*(1e0 - Q3d->p[0][qp] - Q3d->p[1][qp] - Q3d->p[2][qp])
  	   + (*vq)(ind[1])*Q3d->p[0][qp]
  	   + (*vq)(ind[2])*Q3d->p[1][qp]
  	   + (*vq)(ind[3])*Q3d->p[2][qp];

  	vy = (*vq)(offset+ind[0])*(1e0 - Q3d->p[0][qp] - Q3d->p[1][qp] - Q3d->p[2][qp])
  	   + (*vq)(offset+ind[1])*Q3d->p[0][qp]
  	   + (*vq)(offset+ind[2])*Q3d->p[1][qp]
	   + (*vq)(offset+ind[3])*Q3d->p[2][qp];

  	vz = (*vq)(2*offset+ind[0])*(1e0 - Q3d->p[0][qp] - Q3d->p[1][qp] - Q3d->p[2][qp])
  	   + (*vq)(2*offset+ind[1])*Q3d->p[0][qp]
  	   + (*vq)(2*offset+ind[2])*Q3d->p[1][qp]
	   + (*vq)(2*offset+ind[3])*Q3d->p[2][qp];
}


void C3DToFProblem::setMesh(Mesh *m, SaveCells cells)
{
    loadMesh(m->getType(), m->fileName().c_str(), cells);
}


CVar C3DToFProblem::volumeIntegral(int comp, int elem, int qp)
{
	return (U[0]->grad|gphi)*1e-2 + (vx*U[0]->grad(0) + vy*U[0]->grad(1) + vz*U[0]->grad(2))*phi - phi*1;
}




